# coding: utf-8

import sys
sys.path.append("../")
from utils import *
from channel_codec.ldpc_codec import *
import argparse
import cv2
import glob
import time
from tqdm import tqdm
from multiprocessing import Process
import numba



def parse_args():
    parse = argparse.ArgumentParser()
    parse.add_argument(
        "--input_folder", "-if", type=str, required=True,
        help="The input data folder."
    )
    parse.add_argument(
        "--output_folder", "-of", type=str, default="./result",
        help="The output result folder."
    )
    parse.add_argument(
        "--snr", "-snr", type=int, default=5,
        help="The signal to noise you want to set."
    )
    parse.add_argument(
        "--ldpc_rate", "-lr", type=float, default=1/2,
        help="LDPC rate."
    )
    parse.add_argument(
        "--qam_m", "-m", type=int, default=16,
        help="QAM order."
    )
    parse.add_argument(
        "--ldpc_design_file", "-ldf", type=str, default="./channel_codec/1440.720.txt",
        help="ldpc design file path."
    )
    parse.add_argument(
        "--channel_type", "-cht", type=str, default="awgn", choices=["awgn", "rayleigh", "rician"],
        help="The channel type. support [awgn, rayleigh, rician]."
    )
    parse.add_argument(
        "--suffix", "-sf", type=str, default=".mkv", choices=[".mkv", ".png"],
        help="The image format suffix. default .mkv. *.mkv should be one frame picture."
    )
    args = parse.parse_args()
    return args


@numba.jit(forceobj=True, parallel=True)
def multiprocess_one_image(input_path, args):
    input_name = os.path.basename(input_path).split(".")[0]
    if args.suffix == ".mkv":
        video = cv2.VideoCapture(input_path)
        ret, frame = video.read()
        image = Image.fromarray(frame)
    elif args.suffix == ".png":
        image = Image.open(input_path)
    image_bits, shape = image2bits(image)
    # ldpc_decode_bits = ldpc_qam_channel.transmite(image_bits)
    if args.channel_type == "awgn":
        ldpc_decode_bits = ldpc_qam_awgn(image_bits, args.snr, args.qam_m, args.ldpc_design_file)
    elif args.channel_type == "rayleigh":
        ldpc_decode_bits = ldpc_qam_rayleigh(image_bits, args.snr, args.qam_m, args.ldpc_design_file)
    elif args.channel_type == "rician":
        ldpc_decode_bits = ldpc_qam_rician(image_bits, args.snr, args.qam_m, args.ldpc_design_file)
    receive_image = bits2image(ldpc_decode_bits, shape)
    # out_string = ''
    # for i in ldpc_decode_bits:
    #     out_string += str(i)
    # receive_image = bit_to_img(out_string)
    save_image(receive_image, save_folder=args.output_folder, name=input_name, suffix=".png")
    print(f"Image {input_name} saved.")


def _main():
    args = parse_args()
    print(args)
    os.makedirs(args.output_folder, exist_ok=True)
    # ldpc_qam_channel = LdpcQamChannel(
    #     args.output_folder,
    #     args.snr,
    #     args.ldpc_rate,
    #     args.qam_m,
    #     args.ldpc_design_file,
    #     args.channel_type
    # )
    input_paths = glob.glob(os.path.join(args.input_folder, "*.mkv"))
    start = time.time()
    multi_processes = []
    for input_path in tqdm(input_paths):
        # 多进程加速
        # input_name = os.path.basename(input_path).split(".")[0]
        # if args.suffix == ".mkv":
        #     video = cv2.VideoCapture(input_path)
        #     ret, frame = video.read()
        #     image = Image.fromarray(frame)
        # elif args.suffix == ".png":
        #     image = Image.open(input_path)
        # image_bits, shape = image2bits(image)
        # # ldpc_decode_bits = ldpc_qam_channel.transmite(image_bits)
        # if args.channel_type == "awgn":
        #     ldpc_decode_bits = ldpc_qam_awgn(image_bits, args.snr, args.qam_m, args.ldpc_design_file)
        # receive_image = bits2image(ldpc_decode_bits, shape)
        # # out_string = ''
        # # for i in ldpc_decode_bits:
        # #     out_string += str(i)
        # # receive_image = bit_to_img(out_string)
        # save_image(receive_image, input_name, suffix=".png")
        p = Process(target=multiprocess_one_image, args=(input_path, args))
        p.start()
        multi_processes.append(p)
        if len(multi_processes) > 40:
            for cur_process in tqdm(multi_processes):
                cur_process.join()
                multi_processes.pop(0)
        for cur_process in tqdm(multi_processes):
            cur_process.join()
    print(f"Consuming {time.time() - start} secs to process file.")


if __name__ == "__main__":
    _main()
